##Function to calculate the number and percent of thalamocortical connections showing a significant association with an environmental characteristic
sig.effects <- function(df, env.measure){
  enveffects.df <- df
  enveffects.df$significant <- p.adjust(enveffects.df$GAM.cov.pvalue, method = c("fdr")) < 0.05
  sigeffect.totaln <- sum(enveffects.df$significant)
  sigeffect.percent <- round(sigeffect.totaln/length(enveffects.df$significant)*100, 2)
  print(sprintf("%s thalamocortical connections (%s percent) show a significant association between thalamocortical FA and the %s", sigeffect.totaln, sigeffect.percent, env.measure))
}

spin.data <- left_join(spin.parcels, sigeffects.df, by = "tract") #merged data for spin tests, with all 360 parcels in the expected right hemisphere --> left hemisphere order (matching glasser.coords)
spin.data <- merge(spin.data, S.A.axis, by = c("label", "orig_parcelname", "tract"), sort =F)

cor.test(spin.data$GAM.cov.tvalue, spin.data$SA.axis, method = c("spearman"))

perm.sphere.p(x = spin.data$SA.axis, y = spin.data$GAM.cov.tvalue, perm.id = perm.id.full, corr.type = "spearman") 

ggplot() +
  geom_point(data = spin.data, aes(x = SA.axis, y = GAM.cov.tvalue, color = SA.axis), size = 0.5) +
  geom_smooth(data = spin.data, aes(x = SA.axis, y = GAM.cov.tvalue), method = "lm", color = c("black"), formula = y ~ x, linewidth = 0.2) + 
  scale_y_continuous(limits = c(.5, 10)) +
  scale_color_gradient2(low = "goldenrod1", mid = "#ede4f5", high = "#6f1282", guide = "colorbar", aesthetics = "color", midpoint = 180) +
  theme_classic() +
  labs(x="\nS-A axis", y="Environment effect (t)") +
  theme(
  legend.position = "none", 
  axis.text = element_text(size = 5, family = "Arial", color = c("black")),
  axis.title.x = element_text(size = 5, family ="Arial", color = c("black")),
  axis.title.y = element_text(size = 5, family ="Arial", color = c("black")),
  axis.line = element_line(linewidth = .2), 
  axis.ticks = element_line(linewidth = .2)) 

ggsave(filename = "/cbica/projects/thalamocortical_development/figures/images/Figure7/EnvEffect_SArank_correlationplot.pdf", device = "pdf", dpi = 500, width = 1.61, height = 1.3)

#Glasser parcel x,y,z coordinates
load(file = "/cbica/projects/thalamocortical_development/Maps/parcellations/surface/hcp_mmp1.0.rda") #RDA file from the brainGraph package https://github.com/cwatson/brainGraph/blob/master/data/hcp_mmp1.0.rda with glasser region x, y, and z MNI coordinates. regions are in lh --> rh order
hcp_mmp1.0$x.mni[1:180]  <- hcp_mmp1.0$x.mni[1:180]*-1 #convert R-> L hemisphere coords into medial lateral coords
hcp_mmp1.0$z.mni <- hcp_mmp1.0$z.mni*-1 
hcp_mmp1.0$label <- NA
hcp_mmp1.0$label[1:180] <- glasser.regions$label[181:360]
hcp_mmp1.0$label[181:360] <- glasser.regions$label[1:180]
hcp_mmp1.0 <- hcp_mmp1.0 %>% select(-hemi)

#C-Mt files for defining the group-level core-matrix gradient
CPt.glasser.pnc <- read.csv("/cbica/projects/thalamocortical_development/thalamocortical_structuralconnectivity/individual/PNC/PNC_tractaverage_CPt_primary.csv") %>% select("tract", "label", "orig_parcelname", "mean_CPt")
CPt.glasser.hcpd <- read.csv("/cbica/projects/thalamocortical_development/thalamocortical_structuralconnectivity/individual/HCPD/HCPD_tractaverage_CPt_primary.csv")  %>% select("tract", "label", "orig_parcelname", "mean_CPt")
CPt.gradient <- merge(CPt.glasser.pnc, CPt.glasser.hcpd, by = c("label", "tract", "orig_parcelname"), all.x = TRUE, all.y = TRUE, suffixes = c("_pnc", "_hcpd"))
CPt.gradient$CPt <- CPt.gradient %>% select(mean_CPt_pnc, mean_CPt_hcpd) %>% rowMeans(na.rm = T)
CPt.gradient <- CPt.gradient %>% select(CPt, label)

hcp_mmp1.0 <- left_join(hcp_mmp1.0, CPt.gradient, by = "label")

#Function to compute the correlation between an environment map and an anatomical axis (x,y,z) as well as test the significance of the difference between the correlation with the anatomical axis versus the S-A axis
compute_axis_correlation <- function(axis){
  df <- merge(spin.data, hcp_mmp1.0, by = "label", sort = F)
  completeobs.n <- sum(!is.na(df$GAM.cov.tvalue))
  
  #S-A axis - environment correlation (as computed above) for comparison of two overlapping correlations based on dependent groups
  SA.axis.cor <-  cor(df$SA.axis, df$GAM.cov.tvalue, use = "complete.obs", method = c("spearman"))
  
  #Anatomical axis - environment correlation
  anatomical.axis.cor <- cor(df[,axis], df$GAM.cov.tvalue, use = "complete.obs", method = c("spearman")) #correlation between anatomical axis and env metric
  anatomical.axis.pvalue <- perm.sphere.p(x = df[,axis], y = df$GAM.cov.tvalue, perm.id = perm.id.full, corr.type = "spearman") #spin based p-value for the correlation
  
  #S-A axis - anatomical axis correlation
  SA.anatomical.cor <- cor(df$SA.axis, df[,axis], use = "complete.obs", method = c("spearman"))
  
  #Test for significant difference between two correlations based on dependent groups (i.e., correlations with an overlapping input)
  cocor.result <- cocor.pvalue <- cocor.dep.groups.overlap(
  r.jk = SA.axis.cor, 
  r.jh = anatomical.axis.cor, 
  r.kh = SA.anatomical.cor, 
  n = completeobs.n, alternative = "two.sided", test = "hittner2003")
  
  cocor.pvalue <- cocor.result@hittner2003[["p.value"]]
  
  comparison.results <- data.frame(axis = axis, axis.cor = anatomical.axis.cor, axis.cor.pvalue = anatomical.axis.pvalue, SA.cor = SA.axis.cor, cocor.pvalue = cocor.pvalue)
  
  return(comparison.results)
  }

compute_axis_correlation(axis = "y.mni")

compute_axis_correlation(axis = "x.mni")

compute_axis_correlation(axis = "z.mni")

compute_axis_correlation(axis = "CPt")

axis.cor.ps <- c(0.0056, 0.38685, 0.10335, 0.07605, 0.15965) #S-A, A-P, M-L, D-V, C-M
p.adjust(axis.cor.ps, method = c("fdr"))

anatomical.ps <- c(0.0005890574, 0.3319559, 0.3909002, 4.3867e-07) #A-P, M-L, D-V, C-M
p.adjust(anatomical.ps, method = c("fdr"))

#Put above results into df for plotting
axis.results <- data.frame(Axis = factor(), corr = double())
axis.results <- axis.results %>% add_row(Axis = "S-A", corr = 0.3010667)
axis.results <- axis.results %>% add_row(Axis = "A-P", corr = 0.06846274)
axis.results <- axis.results %>% add_row(Axis = "D-V", corr = 0.2117744)
axis.results <- axis.results %>% add_row(Axis = "M-L", corr = 0.1841123)
axis.results <- axis.results %>% add_row(Axis = "C-M", corr = -0.1265908)

axis.results$Axis <- factor(axis.results$Axis, ordered = TRUE, levels = c("S-A", "D-V", "M-L", "A-P", "C-M"))

ggplot(axis.results, aes(x = Axis, y = corr, fill = Axis)) + 
  geom_bar(stat='identity') + 
  labs(x="") +
  labs(y="\nEnvironment Effect Correlation") +
  theme_classic()+
  scale_fill_manual(values=c(alpha("#323280", 0.5), "#672975", "#672975", "#672975", "#672975")) +
  theme(
  legend.position = "none", 
  axis.text = element_text(size = 5, family = "Arial", color = c("black")),
  axis.title.x = element_text(size = 5, family ="Arial", color = c("black")),
  axis.title.y = element_text(size = 5, family ="Arial", color = c("black")),
  axis.line = element_line(linewidth = .2), 
  axis.ticks = element_line(linewidth = .2))

ggsave(filename = "/cbica/projects/thalamocortical_development/figures/images/Figure7/EnvEffect_anatomicalaxes_barplot.pdf", device = "pdf", dpi = 500, width = 1.58, height = 1.3)

#Calculate empirical enrichment statistic for each S-A axis bin
envSES_effects_SAaxisbins <- spin.data %>% 
                             group_by(SA.axis.bin) %>% 
                             do(mean.tvalue.significant = mean(.$GAM.cov.tvalue, na.rm = TRUE)) %>%
                             unnest(cols = c("mean.tvalue.significant"))

envSES_effects_SAaxisbins

#Calculate null distribution of t-values by bin based on spherical spatial permutations (spins)

## inputs to spatially permute (i.e., spin) with the pre-computed spatial permutation matrix 
x <- spin.data$SA.axis.bin #cortical map 1 to spatially permute: S-A axis bins
y <- spin.data$GAM.cov.tvalue #cortical map 2 to spatially permute: significant envSES-thalamocortical FA t-values
perm.id <- perm.id.full

## spin the empirical data 
nroi = dim(perm.id)[1]  #number of regions
nperm = dim(perm.id)[2] #number of permutations
x.perm = y.perm = array(NA,dim=c(nroi,nperm)) #initialize
  for (r in 1:nperm) {
    for (i in 1:nroi) {
      x.perm[i,r] = x[perm.id[i,r]] #spinning x, spatially permuted bins
      y.perm[i,r] = y[perm.id[i,r]] #spinning y, spatially permuted t-values
    }
  }

## compute mean t-value for each S-A axis bin using spatially permuted y inputs 
### set up spun (y) and empirical (x) df
y.spatialperm <- as.data.frame(y.perm) %>% set_names(sprintf("perm%s", seq(from = 1, to = ncol(y.perm))))
y.spatialperm.empiricalx <- cbind(x, y.spatialperm) %>% as.data.frame()
colnames(y.spatialperm.empiricalx)[1] <- c("empirical.x")

### compute mean t-values for S-A axis bins for all spins
envSES_effects_SAaxisbins_permutedy <- y.spatialperm.empiricalx %>% group_by(empirical.x) %>% #group by empirical S-A axis bins
                                       dplyr::summarize(across(everything(), \(x) mean(x, na.rm = TRUE))) %>% #calculate mean t-value for permuted data
                                       select(-empirical.x) %>% t() %>% as.data.frame() %>%
                                       set_names(c("SA.bin1.tvalue","SA.bin2.tvalue", "SA.bin3.tvalue", "SA.bin4.tvalue", "SA.bin5.tvalue"))

#Function to plot the null distribution of environment effect t-values for a given S-A axis bin and test whether the true t-value is significantly smaller or greater in magnitude than the null (permutation based p)
enveffects_SAbin_enrichment <- function(bin, enrichment_type){
  
  nperm = 10000
  
  SAbin.tvalue.empirical <- envSES_effects_SAaxisbins %>% filter(SA.axis.bin == bin) %>% select(mean.tvalue.significant)
  
  #Histogram
  SA.colorvector <- c("#FFC228", "#FFD16A", "#e9dcf2", "#BE93C4", "#6F1382")
  
  null.tvalue.histogram <- envSES_effects_SAaxisbins_permutedy %>% select(sprintf("SA.bin%s.tvalue", bin)) %>%
    ggplot(aes(x = .[,1])) + geom_histogram(fill = c(SA.colorvector)[bin]) +
    theme_classic() +
    geom_vline(xintercept =  SAbin.tvalue.empirical$mean.tvalue.significant, linewidth = 0.25) +
    xlab("Environment effect (t-value)") +
    ylab("Number of nulls") +
    scale_x_continuous(breaks = c(2, 3, 4, 5), limits = c(1.25, 5.25)) +
    theme(
     legend.position = "none", 
     axis.text.x = element_text(size = 5, family = "Arial", color = c("black")),
     axis.text.y = element_blank(),
     axis.title.x = element_text(size = 5, family ="Arial", color = c("black")),
     axis.title.y = element_text(size = 5, family ="Arial", color = c("black")),
     axis.line.x = element_line(linewidth = .2), 
     axis.line.y = element_blank(),
     axis.ticks = element_line(linewidth = .2))
    
  #Permutation-based p
    if(enrichment_type == "smaller"){
      perm.p <- sum(envSES_effects_SAaxisbins_permutedy[,bin] < SAbin.tvalue.empirical$mean.tvalue.significant)/nperm
    }
  
  if(enrichment_type == "larger"){
      perm.p <- sum(envSES_effects_SAaxisbins_permutedy[,bin] > SAbin.tvalue.empirical$mean.tvalue.significant)/nperm
  }
  
  permutation.results <- list(null.tvalue.histogram, perm.p)
  return(permutation.results)
}

enveffects_SAbin_enrichment(bin = 1, enrichment_type = "smaller")

ggsave(filename = "/cbica/projects/thalamocortical_development/figures/images/Figure7/quintile1_enveffect_spins.pdf", device = "pdf", dpi = 500, width = .95, height =.9)

enveffects_SAbin_enrichment(bin = 2, enrichment_type = "smaller")

ggsave(filename = "/cbica/projects/thalamocortical_development/figures/images/Figure7/quintile2_enveffect_spins.pdf", device = "pdf", dpi = 500, width = .95, height =.9)

enveffects_SAbin_enrichment(bin = 3, enrichment_type = "smaller")

ggsave(filename = "/cbica/projects/thalamocortical_development/figures/images/Figure7/quintile3_enveffect_spins.pdf", device = "pdf", dpi = 500, width = .95, height =.9)

enveffects_SAbin_enrichment(bin = 4, enrichment_type = "larger")

ggsave(filename = "/cbica/projects/thalamocortical_development/figures/images/Figure7/quintile4_enveffect_spins.pdf", device = "pdf", dpi = 500, width = .95, height =.9)

enveffects_SAbin_enrichment(bin = 5, enrichment_type = "larger")

ggsave(filename = "/cbica/projects/thalamocortical_development/figures/images/Figure7/quintile5_enveffect_spins.pdf", device = "pdf", dpi = 500, width = .95, height =.9)